from ultralytics import YOLO
import pandas as pd
import os
 
TubeAndLevel = r"D:\Desktop\YOLO-RotatedBarocde\model\tube\last.pt" #试管和液位识别模型
directory = r"D:\Desktop\大论文\章节4-试管检测和条码识别系统实现\试管液位图片" #存放图片路径
save_dir = r"D:\Desktop\大论文\章节4-试管检测和条码识别系统实现\预测结果"

class TestTubeDetector:
    def __init__(self,tube_model,barcode_model,source,save_dir):
        self.tube_model = YOLO(tube_model) #试管模型
        
        self.source = source #处理图片路径
        
        self.level_percent = []
        self.cap_color = []
        self.image_list = []
        self.color_map = {
            3 : 'red',
            4 : 'orange',
            5 : 'yellow',
            6 : 'green',
            7 : 'blue',
            8 : 'purple',
            9 : 'black',
            10 : 'grey',
            11 : 'cover',
            12 : 'greenplug'
        }
        self.project = save_dir

    def detect_tube(self):
        #获取路径下的所有图片
        for file in os.listdir(self.source):
            file_path = os.path.join(self.source,file)
            metric = self.tube_model.predict(source=file_path) 
            cls = metric[0].boxes.cls.tolist()
            boxes = metric[0].boxes.xyxy.tolist()
            level_center = self.get_level_center_x(cls,boxes)
            tube_left_x,tube_right_x = self.get_tube_border_x(cls,boxes)
            
            #液位百分比和管帽颜色
            percent = self.get_level_percent(level_center,tube_left_x,tube_right_x)
            color =  self.get_tube_cap_color(cls)
            self.level_percent.append(percent)
            self.cap_color.append(color)
            self.image_list.append(file)
            
    def get_level_center_x(self,cls,boxes):
        for index in range(len(cls)):
            if int(cls[index]) == 1: #level的id为1
                return (boxes[index][0] + boxes[index][2]) / 2
        return None
    
    def get_tube_border_x(self,cls,boxes):
        for index in range(len(cls)):
            if int(cls[index]) == 0 or int(cls[index]) == 2:
                return boxes[index][0],boxes[index][2]
        return None
    
    def get_level_percent(self,level_center,left,right):
        tube_weight = right - left
        level_weight = level_center - left
        return level_weight / tube_weight
    
    def get_tube_cap_color(self,cls):
        for index in range(len(cls)):
            if int(cls[index]) > 2:
                return self.color_map[cls[index]]
                
    def save_to_excel(self):
        df = pd.DataFrame({
            "name" : self.image_list,
            "color" : self.cap_color,
            "percent" : self.level_percent
        })
        
        #保存位置
        file_name = os.path.join(self.project,'result.xlsx')
        df.to_excel(file_name)
        
detector = TestTubeDetector(TubeAndLevel,None,directory,save_dir)
detector.detect_tube()
detector.save_to_excel()